"""
Factory module for backbone creation.

This module provides factory functions for creating backbone models
with appropriate configuration handling and weight loading.
"""

import logging
import os
from typing import Dict

import torch
from torch import nn

from models.config import ModelConfig, BackboneConfig, ViTConfig
from models.interfaces import BackboneBase
from models import BACKBONE_REGISTRY  # 从models包导入

logger = logging.getLogger(__name__)


class BackboneFactory:
    """创建backbone模型的工厂类。"""

    @staticmethod
    def create(config: ModelConfig) -> BackboneBase:
        """Create a backbone model based on configuration.

        Args:
            config: Model configuration object

        Returns:
            Instantiated backbone model

        Raises:
            ValueError: If backbone_type is not registered
            RuntimeError: If weight loading fails
        """
        backbone_type = config.backbone_type
        backbone_config = config.backbone

        # 使用正确的注册表
        if backbone_type not in BACKBONE_REGISTRY:
            available_backbones = BACKBONE_REGISTRY.list_registered()  # 方法名可能不同
            raise ValueError(
                f"Backbone type {backbone_type} is not registered. "
                f"Available backbones: {available_backbones}"
            )

        logger.info(f"Creating backbone: {backbone_type}")

        # Create the backbone model
        backbone_fn = BACKBONE_REGISTRY.get(backbone_type)
        backbone = backbone_fn(backbone_config)

        # Return the backbone
        return backbone

    @staticmethod
    def load_pretrained(
        backbone: BackboneBase, pretrained_path: str, strict: bool = False
    ) -> BackboneBase:
        """Load pretrained weights into backbone.

        Args:
            backbone: The backbone model to load weights into
            pretrained_path: Path to pretrained weights
            strict: Whether to strictly enforce that the keys in state_dict match

        Returns:
            Backbone with loaded weights

        Raises:
            FileNotFoundError: If pretrained_path does not exist
            RuntimeError: If weight loading fails
        """
        if not os.path.exists(pretrained_path):
            raise FileNotFoundError(f"Pretrained model not found at {pretrained_path}")

        logger.info(f"Loading pretrained weights from {pretrained_path}")

        # Load state dict
        state_dict = torch.load(pretrained_path, map_location="cpu")

        # Handle different state dict formats
        if "model" in state_dict:
            state_dict = state_dict["model"]
        elif "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]

        # Filter out unexpected keys
        model_state_dict = backbone.state_dict()
        filtered_state_dict = {
            k: v for k, v in state_dict.items() if k in model_state_dict
        }

        # Report missing and unexpected keys
        missing_keys = [
            k for k in model_state_dict.keys() if k not in filtered_state_dict
        ]
        unexpected_keys = [k for k in state_dict.keys() if k not in model_state_dict]

        if missing_keys:
            logger.warning(f"Missing keys in pretrained weights: {missing_keys}")
        if unexpected_keys:
            logger.info(f"Unexpected keys in pretrained weights: {unexpected_keys}")

        # Load the weights
        backbone.load_state_dict(filtered_state_dict, strict=strict)

        return backbone

    @staticmethod
    def get_feature_dims(backbone: BackboneBase) -> Dict[str, int]:
        """Get feature dimensions from backbone.

        Args:
            backbone: The backbone model

        Returns:
            Dictionary mapping feature names to dimensions
        """
        # For now, just return the output dimension, can be extended for multi-level features
        return {"output": backbone.get_output_dim()}
